"""
LangGraph工作流可视化工具
用于生成工作流流程图
"""

import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.patches import FancyBboxPatch, ConnectionPatch
import numpy as np
from typing import Dict, List, Tuple
import os

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

class WorkflowVisualizer:
    """工作流可视化器"""
    
    def __init__(self):
        """初始化可视化器"""
        self.fig = None
        self.ax = None
        self.node_positions = {}
        self.node_colors = {
            'start': '#4CAF50',      # 绿色 - 开始节点
            'process': '#2196F3',    # 蓝色 - 处理节点
            'decision': '#FF9800',   # 橙色 - 决策节点
            'end': '#F44336',        # 红色 - 结束节点
            'tool': '#9C27B0'        # 紫色 - 工具节点
        }
    
    def create_workflow_diagram(self, output_path: str = "travel_agent_workflow.png"):
        """创建旅游助手工作流图"""
        # 创建图形
        self.fig, self.ax = plt.subplots(1, 1, figsize=(16, 12))
        self.ax.set_xlim(0, 10)
        self.ax.set_ylim(0, 10)
        self.ax.axis('off')
        
        # 定义节点位置
        self._define_node_positions()
        
        # 绘制节点
        self._draw_nodes()
        
        # 绘制连接线
        self._draw_connections()
        
        # 添加标题和说明
        self._add_title_and_legend()
        
        # 保存图片
        plt.tight_layout()
        plt.savefig(output_path, dpi=300, bbox_inches='tight', 
                   facecolor='white', edgecolor='none')
        plt.close()
        
        print(f"工作流图已保存到: {output_path}")
    
    def _define_node_positions(self):
        """定义节点位置"""
        self.node_positions = {
            # 主流程节点
            'start': (5, 9.5),
            'classify_query': (5, 8.5),
            
            # 分支节点
            'handle_travel_guide': (2, 7),
            'extract_route_info': (8, 7),
            'handle_route_planning': (8, 6),
            
            # 路线规划子流程
            'plan_route': (8, 5),
            'ask_clarification': (8, 4),
            
            # 检查和处理节点
            'check_relevance': (5, 4),
            'generate_final_answer': (5, 3),
            'end': (5, 2),
            
            # 工具节点
            'rag_tool': (1, 6),
            'amap_tool': (9, 4.5)
        }
    
    def _draw_nodes(self):
        """绘制节点"""
        # 开始节点
        self._draw_node('start', '开始', 'start')
        
        # 查询分类节点
        self._draw_node('classify_query', '查询分类\n(旅游攻略/路线规划)', 'decision')
        
        # 旅游攻略处理节点
        self._draw_node('handle_travel_guide', '处理旅游攻略查询', 'process')
        self._draw_node('rag_tool', 'RAG工具\n检索文档', 'tool')
        
        # 路线规划处理节点
        self._draw_node('extract_route_info', '提取路线信息\n(出发地/目的地)', 'process')
        self._draw_node('handle_route_planning', '处理路线规划', 'process')
        self._draw_node('plan_route', '规划路线', 'process')
        self._draw_node('amap_tool', '高德地图工具\n路线规划', 'tool')
        self._draw_node('ask_clarification', '询问澄清信息', 'process')
        
        # 检查和生成节点
        self._draw_node('check_relevance', '检查答案相关性', 'decision')
        self._draw_node('generate_final_answer', '生成最终答案', 'process')
        
        # 结束节点
        self._draw_node('end', '结束', 'end')
    
    def _draw_node(self, node_id: str, label: str, node_type: str):
        """绘制单个节点"""
        x, y = self.node_positions[node_id]
        color = self.node_colors[node_type]
        
        # 根据节点类型选择形状
        if node_type == 'decision':
            # 菱形
            width, height = 1.2, 0.8
            bbox = FancyBboxPatch((x - width/2, y - height/2), width, height,
                                boxstyle="round,pad=0.1", 
                                facecolor=color, edgecolor='black', linewidth=1.5)
        elif node_type == 'tool':
            # 圆角矩形
            width, height = 1.4, 0.6
            bbox = FancyBboxPatch((x - width/2, y - height/2), width, height,
                                boxstyle="round,pad=0.1", 
                                facecolor=color, edgecolor='black', linewidth=1.5)
        else:
            # 矩形
            width, height = 1.4, 0.6
            bbox = FancyBboxPatch((x - width/2, y - height/2), width, height,
                                boxstyle="round,pad=0.1", 
                                facecolor=color, edgecolor='black', linewidth=1.5)
        
        self.ax.add_patch(bbox)
        
        # 添加文本
        self.ax.text(x, y, label, ha='center', va='center', 
                    fontsize=9, fontweight='bold', color='white')
    
    def _draw_connections(self):
        """绘制连接线"""
        connections = [
            # 主流程
            ('start', 'classify_query'),
            ('classify_query', 'handle_travel_guide', '旅游攻略'),
            ('classify_query', 'extract_route_info', '路线规划'),
            ('classify_query', 'check_relevance', '未知类型'),
            
            # 旅游攻略分支
            ('handle_travel_guide', 'rag_tool'),
            ('rag_tool', 'check_relevance'),
            
            # 路线规划分支
            ('extract_route_info', 'handle_route_planning'),
            ('extract_route_info', 'ask_clarification', '信息不完整'),
            ('handle_route_planning', 'plan_route'),
            ('plan_route', 'amap_tool'),
            ('amap_tool', 'check_relevance'),
            
            # 最终处理
            ('check_relevance', 'generate_final_answer'),
            ('generate_final_answer', 'end'),
            ('ask_clarification', 'end')
        ]
        
        for connection in connections:
            if len(connection) == 2:
                start_node, end_node = connection
                label = None
            else:
                start_node, end_node, label = connection
            
            self._draw_arrow(start_node, end_node, label)
    
    def _draw_arrow(self, start_node: str, end_node: str, label: str = None):
        """绘制箭头连接"""
        start_pos = self.node_positions[start_node]
        end_pos = self.node_positions[end_node]
        
        # 计算箭头位置（避免重叠）
        dx = end_pos[0] - start_pos[0]
        dy = end_pos[1] - start_pos[1]
        distance = np.sqrt(dx**2 + dy**2)
        
        # 调整起点和终点，避免与节点重叠
        start_x = start_pos[0] + (dx / distance) * 0.3
        start_y = start_pos[1] + (dy / distance) * 0.3
        end_x = end_pos[0] - (dx / distance) * 0.3
        end_y = end_pos[1] - (dy / distance) * 0.3
        
        # 绘制箭头
        arrow = mpatches.FancyArrowPatch(
            (start_x, start_y), (end_x, end_y),
            arrowstyle='->', mutation_scale=20,
            color='black', linewidth=1.5
        )
        self.ax.add_patch(arrow)
        
        # 添加标签
        if label:
            mid_x = (start_x + end_x) / 2
            mid_y = (start_y + end_y) / 2
            
            # 调整标签位置，避免与箭头重叠
            if abs(dy) > abs(dx):  # 垂直箭头
                mid_x += 0.3 if dx > 0 else -0.3
            else:  # 水平箭头
                mid_y += 0.2 if dy > 0 else -0.2
            
            self.ax.text(mid_x, mid_y, label, ha='center', va='center',
                        fontsize=8, bbox=dict(boxstyle="round,pad=0.2", 
                        facecolor='white', edgecolor='gray', alpha=0.8))
    
    def _add_title_and_legend(self):
        """添加标题和图例"""
        # 标题
        self.ax.text(5, 10.2, '智能旅游助手 Agent 工作流', 
                    ha='center', va='center', fontsize=16, fontweight='bold')
        
        # 图例
        legend_elements = [
            mpatches.Patch(color=self.node_colors['start'], label='开始节点'),
            mpatches.Patch(color=self.node_colors['process'], label='处理节点'),
            mpatches.Patch(color=self.node_colors['decision'], label='决策节点'),
            mpatches.Patch(color=self.node_colors['tool'], label='工具节点'),
            mpatches.Patch(color=self.node_colors['end'], label='结束节点')
        ]
        
        self.ax.legend(handles=legend_elements, loc='upper right', 
                      bbox_to_anchor=(0.98, 0.98), fontsize=10)
        
        # 添加说明文本
        description = """
工作流说明：
1. 用户输入查询后，系统首先进行分类判断
2. 旅游攻略查询：调用RAG工具检索相关文档
3. 路线规划查询：提取出发地和目的地信息，调用高德地图工具
4. 信息不完整时：询问用户澄清
5. 最终检查答案相关性并生成回复
        """
        
        self.ax.text(0.5, 1.5, description, ha='left', va='top', 
                    fontsize=9, bbox=dict(boxstyle="round,pad=0.5", 
                    facecolor='lightblue', edgecolor='blue', alpha=0.8))

def main():
    """主函数"""
    visualizer = WorkflowVisualizer()
    
    # 创建输出目录
    output_dir = "/Users/pengxiong/workspace/pyWorkspace/traval-agent"
    output_path = os.path.join(output_dir, "travel_agent_workflow.png")
    
    # 生成工作流图
    visualizer.create_workflow_diagram(output_path)

if __name__ == "__main__":
    main()
